

import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;
//import java.util.TreeSet;
import java.util.Map.Entry;
import java.net.*;
import java.io.*;

public class Node {

	public enum State {
		INITWAIT, INITCONNECT, CHANGELINKCOST, WAITLINKCOSTCHANGE,
		WAITFORNEIGHBORS
	}
	
	// map of ID --> Port#
	HashMap<Integer, Integer> MapIDPort = new HashMap<Integer, Integer>();
	// map of port# --> ID
	HashMap<Integer, Integer> MapPortID = new HashMap<Integer, Integer>();
	
	Set<Integer> _neighborsPorts = new HashSet<Integer>();		// neighbors' port numbers
	DatagramSocket _udpSocket;
	DatagramPacket _inPacket;
	DatagramPacket _outPacket;
	byte[] inBuffer;
	byte[] outBuffer;
	
	int _myport;
	int _myID;			// based on running order
	int _totalNodes;		// total nodes in network
	int MAX = 10000;
	int BUFFERSIZE = 64;		
	int NEWCOST = 30;		// new link cost value
	
	boolean linkCostChanged = false;
	boolean poisonedReverse = true;		// whether or not to use Poisoned Reverse
	
	int[][] DVtable, DVtableCopy;
	int[][] linkCost;
	int[][] via, viaCopy;			// node on path from i to j
									// if i to j is direct then via=j
	
	int[][] via3 = {
			{0, 1, 1},
			{0, 1, 2},
			{1, 1, 2}
	};
	
	int[][] via4 = {
			{0, 1, 1, 1},
			{0, 1, 3, 3},
			{3, 3, 2, 3},
			{1, 1, 2, 3}
	};
	
	int[][] DVtable3 = {
			{0, 4, 5},
			{4, 0, 1},
			{5, 1, 0}
	};
	
	int[][] DVtable4 = {
			{0, 4, 7, 5},
			{4, 0, 3, 1},
			{7, 3, 0, 2},
			{5, 1, 2, 0}
	};
	
	// hard-coded link cost tables for 3 and 4 nodes
	int[][] linkCost3 = {
			{0,  4, 20},
			{4,  0, 1},
			{20, 1, 0}
	};
	
	int[][] linkCost4 = {{0,  4, MAX, 20},
						 {4,  0, 8,   1},
						{MAX, 8, 0,   2},
						{20,  1, 2,   0}};
	
	int[] _neighbors3 = {2, 2, 2};
	int[] _neighbors4 = {2, 3, 2, 3};
		
	State _state;
	
	
	
	public Node(String[] args) {
		poisonedReverse = (args[0].toUpperCase()).equals("T") ? true : false;		
		_totalNodes = Integer.parseInt(args[1]);
		_myID = Integer.parseInt(args[2]);
		_myport = Integer.parseInt(args[3]);
		_state  = State.INITWAIT;
		
		if (args.length > 4) {			
			for (int i=4; i<args.length; i++) {
				_neighborsPorts.add(Integer.parseInt(args[i]));				
			}			
			_state = State.INITCONNECT;					
		}
		try {
			_udpSocket = new DatagramSocket(_myport, InetAddress.getLocalHost());
		} catch (SocketException e) {
			e.printStackTrace();
		} catch (UnknownHostException e) {
			e.printStackTrace();
		}
		
		if (_totalNodes == 3) {
			DVtable = DVtable3;		
			linkCost = linkCost3;	
			via = via3;			
		}
		else {
			DVtable = DVtable4;
			linkCost = linkCost4;
			via = via4;
		}
		DVtableCopy = DVtable;
		viaCopy = via;
		
		System.out.println("ID: " + _myID + " UDP Port: " + _myport);
	}	
	
	public void Start() throws InterruptedException {
		while (true) {
			switch (_state) {
			case INITWAIT:
				initWaitForConnection();
				break;		
			case INITCONNECT:
				initConnect();
				break;
			case CHANGELINKCOST:
				changeLinkCost();
				break;
			case WAITLINKCOSTCHANGE:
				//nothing
				break;			
			case WAITFORNEIGHBORS:
				waitForNeighbors();
				break;
			}
			try {
				Thread.sleep(10);
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
		}
	}
	
	
	/**
	 * Wait to receive DV from neighbors.
	 */
	private void waitForNeighbors() {
		
		System.out.println("Poisoned Reverse: " + poisonedReverse);
		
		if (linkCostChanged) {
			linkCostChanged = false;
			
			// Update DV table and send update to neighbors
			if (updateDVtable(poisonedReverse) == true) {
				informNeighbors();
			}
			displayDVtable();
		}
		
		System.out.println("Waiting for DV from neighbors....");
	
		// waits until there's a DV from a neighbor
		_inPacket = receivePacket();
		
		int sendingPort = _inPacket.getPort();
		int neighborID = MapPortID.get(sendingPort);
		byte[] data = _inPacket.getData();
		

		// extract data containing both DV and via vector 
		if (data.length >= (4*_totalNodes*2)) {
			// update my DV table
			copyDataToDVtable(neighborID, data);
			
			//displayDVtable();
			
			if (updateDVtable(poisonedReverse) == true) {
				informNeighbors();
			}
			displayDVtable();
			//_state = State.WAITFORNEIGHBORS;
		}
	}
	

	/**
	 * Informs neighbors of changes in my DV. 
	 */
	private void informNeighbors() {

		byte[] myDV = getMyDV();
		for (int neighborPort:_neighborsPorts) {
			// Sends my updated DV to a neighbor's port#.
			sendPacket(neighborPort, myDV);					
		}		
	}

	/**
	 * Displays DV table.
	 */
	private void displayDVtable() {
		
		System.out.println("--------DV table--------");
		System.out.print("\t");
		
		for (int i=0; i<_totalNodes; i++) {
			System.out.print(i + "\t");
		}
		System.out.println("");System.out.println("");
		
		for (int i=0; i<DVtable.length; i++) {
			System.out.print(i + "\t");
			for (int j=0; j<DVtable[i].length; j++) {
				System.out.print(DVtable[i][j] + "\t");
			}
			System.out.println();
		}
		System.out.println("----------------");
	}

	/**
	 * Generates my DV as a byte[].
	 */
	private byte[] getMyDV() {
		byte[] myDV = new byte[4*_totalNodes*2];
		byte[] temp;
		
		for (int i=0; i<_totalNodes; i++) {
			temp = intToByteArray(DVtable[_myID][i]);
			myDV[4*i] = temp[0];
			myDV[4*i+1] = temp[1];
			myDV[4*i+2] = temp[2];
			myDV[4*i+3] = temp[3];			
			
			temp = intToByteArray(via[_myID][i]);
			myDV[4*(i+_totalNodes)] = temp[0];
			myDV[4*(i+_totalNodes)+1] = temp[1];
			myDV[4*(i+_totalNodes)+2] = temp[2];
			myDV[4*(i+_totalNodes)+3] = temp[3];	
		}
		return myDV;
	}


	/**
	 * Copies byte[] in datagram packet to DV table.
	 */
	private void copyDataToDVtable(int neighborID, byte[] data) {
		for (int i=0; i<_totalNodes; i++) {
			DVtable[neighborID][i] = byteArrayToInt(data, 4*i);
			via[neighborID][i] = byteArrayToInt(data, 4*(i+_totalNodes));
		}		
		
		// Display DV received from neighborID
		System.out.println("Received from node " + neighborID);
		System.out.print( neighborID + ": \t");
		
		for (int i=0; i<_totalNodes; i++) {
			System.out.print( DVtable[neighborID][i] + "\t" );
		}	
		System.out.println();
	}

	
	/**
	 * Changes link cost between node 0 and 1(hard-coded) to a new hard-coded value.
	 */
	private void changeLinkCost() {
		
		System.out.println("Updating link cost...");
		
		displayDVtable();		// old DV table
		
		if (_totalNodes == 3) {
			if ((_myID == 0) || (_myID==1)) {
				linkCost[0][1] = NEWCOST;
				linkCost[1][0] = NEWCOST;
				linkCostChanged = true;				
				System.out.println("New link cost between [0] and [1]: " + NEWCOST);
			}
		}
		else {
			// 4 nodes case
			if ((_myID == 0) || (_myID==1)) {
				linkCost[0][1] = NEWCOST;
				linkCost[1][0] = NEWCOST;
				linkCostChanged = true;
				System.out.println("New link cost between [0] and [1]: " + NEWCOST);
			}
		}		
		_state = State.WAITFORNEIGHBORS;
	}
	
	
	/**
	 * Updates a node's Distance vector.
	 */
	private boolean updateDVtable(boolean poisonedReverse) {
		int temp;		
		boolean DVchanged = false;
		
		for (int dest=0; dest<_totalNodes; dest++) {
			int min = MAX;
			if (_myID != dest) {				
				for (int neighbor:_neighborsPorts) {
					
					int nID = MapPortID.get(neighbor);
					if (!poisonedReverse) {										
						temp = linkCost[_myID][nID] + DVtable[nID][dest]; 									
					}
					else { /* using poisoned reverse */
						int pvCost;
						// this node is on the path from neighbor to dest
						// return infinite cost instead
						if (via[nID][dest] == _myID) {
							pvCost = MAX;
						}
						else {
							pvCost = DVtable[nID][dest];
						}
						temp = linkCost[_myID][nID] + pvCost; 						
					}
					if (temp > 0 && temp < min) {						
						min = temp;
						via[_myID][dest] = neighbor;
					}
				}
				// Distance vector has changed
				if (min != DVtable[_myID][dest]) {					
					DVtable[_myID][dest] = min;					
					DVchanged = true;
				}
			}
		}
		
		return DVchanged;		
	}


	/**
	 * Makes initial connections to specified ports.
	 */
	private void initConnect() throws InterruptedException {
				
		for (int neighborPort : _neighborsPorts ) {
			
			//System.out.println( "neighbor port: " + neighborPort );
			// send my ID to neighbor nodes
			sendPacket(neighborPort, intToByteArray(_myID));
			
			// receive packet from neighbor, 
			// map its ID with its port# 
			// do not need to put its port # in neighbor set			
			System.out.println("Waiting for connections...");
			//Thread.sleep(10);
			
			_inPacket = receivePacket();
			int neighborID = byteArrayToInt(_inPacket.getData(), 0);
			//System.out.println( "received port: " + _inPacket.getPort() );
			//System.out.println( "neighbor ID: " + neighborID );
			
			MapIDPort.put(neighborID, _inPacket.getPort());			
		}
		_state = State.INITWAIT;
	}

	/**
	 * Waits until initial connections are established.
	 * @throws IOException
	 */
	void initWaitForConnection() {
		
		if ( _totalNodes == 3) {		
			while (_neighborsPorts.size() < _neighbors3[_myID]) {
				waitForConnection();				
			}
			_state = State.CHANGELINKCOST;
		}
		else {
			while (_neighborsPorts.size() < _neighbors4[_myID]) {
				waitForConnection();
			}
			_state = State.CHANGELINKCOST;
		}
		
		System.out.println( "ID" + "\t" + "Port#" );
		
		for (Entry<Integer, Integer> entry:MapIDPort.entrySet()) {
			MapPortID.put(entry.getValue(), entry.getKey());
			
			System.out.println( entry.getKey() + "\t" + entry.getValue() );
		}
		System.out.println();
	}
	  
	/**
	 * 
	 */
	private void waitForConnection(){
		// receive packet from neighbor, 
		// add port # to my neighbor set
		// also map ID with port#
		System.out.println("Waiting for connections...");
		_inPacket = receivePacket();
		System.out.println("Done.");
		
		int neighborPort = _inPacket.getPort();
		_neighborsPorts.add(neighborPort);
		
		int neighborID = byteArrayToInt(_inPacket.getData(), 0); 
		//System.out.println("received ID:" + neighborID);
		MapIDPort.put(neighborID, neighborPort);
		
		try {
			Thread.sleep(10);
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
		// send my ID to neighbor nodes
		sendPacket(neighborPort, intToByteArray(_myID));		
	}

	/**
	 * Sends a byte buffer to a specified port. 
	 */
	private void sendPacket(int neighborPort, byte[] outBuffer) {

		try {
			_outPacket = new DatagramPacket( outBuffer, outBuffer.length,
											InetAddress.getLocalHost(), neighborPort );
			//System.out.println("sending...");
			_udpSocket.send(_outPacket);
			//System.out.println("sent.");
		} catch (UnknownHostException e) {
			e.printStackTrace();
		}
		catch (IOException e) {
			e.printStackTrace();
		}		
	}
	
	/**
	 * 
	 */
	private DatagramPacket receivePacket() {
		DatagramPacket inPacket = null;
		inBuffer = new byte[BUFFERSIZE];
		inPacket = new DatagramPacket(inBuffer, inBuffer.length);
		try {			
			_udpSocket.receive(inPacket);
			//System.out.println("received.");
		} catch (IOException e) {
			e.printStackTrace();
		}
		return inPacket;
	}
	
	
	/**
	 * Returns a 4-byte array from an int.
	 */
	static byte[] intToByteArray(int value) {
		return new byte[]{  (byte)(value >>> 24),
							(byte)(value >> 16 & 0xff),
							(byte)(value >> 8 & 0xff),
							(byte)(value & 0xff) };			
	}
	
	/**
	 * Returns an int from a 4-byte array.
	 */
	static  int byteArrayToInt(byte[] b, int startIndex) {
        return (b[startIndex] << 24)
                + ((b[startIndex+1] & 0xFF) << 16)
                + ((b[startIndex+2] & 0xFF) << 8)
                + (b[startIndex+3] & 0xFF);
	}
}

